import numpy as np
import matplotlib.pyplot as plt
import os
from pathlib import Path
import matplotlib.colors as mcolors
from itertools import cycle

# Define paths
base_dir = Path('/home/hjx-linux/desktop/spack/opt/spack/linux-ubuntu22.04-skylake/gcc-11.4.0/openfoam-2406-mrfhxvxnx4fv3uhyihxd4f7fxgt2hqd7/HunterPhd/HUAWEI/hpcTest')
data_dir = base_dir / 'python' / 'data'
pdf_dir = base_dir / 'python' / 'pythonPDF'
svg_dir = base_dir / 'python' / 'pythonSVG'

# Create output directories if they don't exist
pdf_dir.mkdir(parents=True, exist_ok=True)
svg_dir.mkdir(parents=True, exist_ok=True)

# Test cases
test_cases = ['advevtionEqn', 'cavity', 'poissonEqn']


# 自适应颜色和标记分配函数
def get_color_marker_mapping(platform_names):
    """
    为给定的平台名称列表自动分配颜色和标记符号
    
    Args:
        platform_names: 平台名称列表
    
    Returns:
        colors_dict: 颜色映射字典
        markers_dict: 标记符号映射字典
    """
    # 使用 matplotlib 的高质量颜色系列
    color_palette = [
        'tab:blue', 'tab:orange', 'tab:green', 'tab:red', 'tab:purple',
        'tab:brown', 'tab:pink', 'tab:gray', 'tab:olive', 'tab:cyan',
        '#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd',
        '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf'
    ]
    
    # 标记符号列表
    marker_list = ['o', 's', '^', 'D', 'v', '<', '>', 'p', '*', 'h', 
                   'H', '+', 'x', 'd', '|', '_']
    
    # 创建循环迭代器，支持任意数量的平台
    color_cycle = cycle(color_palette)
    marker_cycle = cycle(marker_list)
    
    # 为每个平台分配颜色和标记
    colors_dict = {}
    markers_dict = {}
    
    for platform in sorted(platform_names):  # 排序保证一致性
        colors_dict[platform] = next(color_cycle)
        markers_dict[platform] = next(marker_cycle)
    
    return colors_dict, markers_dict

for test_case in test_cases:
    print(f"Processing {test_case}...")
    
    # Create figure with two y-axes
    fig, ax1 = plt.subplots(figsize=(8, 6))
    ax2 = ax1.twinx()
    
    # Read data files
    case_dir = data_dir / test_case
    data_files = list(case_dir.glob('*.txt'))
    
    # 获取所有平台名称并自动分配颜色和标记
    platform_names = [f.stem for f in data_files]
    colors, markers = get_color_marker_mapping(platform_names)
    
    print(f"  Found platforms: {', '.join(sorted(platform_names))}")
    
    # Store all cores for ideal speedup line
    all_cores = set()
    
    for data_file in data_files:
        platform_name = data_file.stem  # Get filename without extension
        
        # Read data
        try:
            data = np.loadtxt(data_file)
            if data.size == 0:
                print(f"Warning: {data_file} is empty, skipping...")
                continue
            
            # Handle single line data
            if data.ndim == 1:
                data = data.reshape(1, -1)
            
            cores = data[:, 0].astype(int)
            times = data[:, 1]
            
            # Calculate speedup (base time / current time)
            base_time = times[0]
            speedup = base_time / times
            
            # Add cores to set for ideal speedup
            all_cores.update(cores)
            
            # Plot time on left axis
            line1 = ax1.plot(cores, times, marker=markers[platform_name], 
                            color=colors[platform_name], 
                            label=f'{platform_name} (time)', linewidth=2, markersize=8)
            
            # Plot speedup on right axis
            line2 = ax2.plot(cores, speedup, marker=markers[platform_name], 
                            color=colors[platform_name], linestyle='--',
                            label=f'{platform_name} (speedup)', linewidth=2, markersize=8, alpha=0.7)
            
        except Exception as e:
            print(f"Error reading {data_file}: {e}")
            continue
    
    # Plot ideal speedup line
    if all_cores:
        sorted_cores = sorted(all_cores)
        min_cores = sorted_cores[0]
        ideal_speedup = [c / min_cores for c in sorted_cores]
        ax2.plot(sorted_cores, ideal_speedup, 'k:', linewidth=2, label='Ideal speedup')
    
    # Set x-axis to log scale
    ax1.set_xscale('log')
    ax2.set_xscale('log')
    
    # Set y-axis (right) to log scale for speedup
    ax2.set_yscale('log')
    
    # Set x-axis ticks to show actual core numbers
    if all_cores:
        sorted_cores = sorted(all_cores)
        ax1.set_xticks(sorted_cores)
        ax1.set_xticklabels(sorted_cores)
        
        # Set right y-axis ticks to show ideal speedup values
        min_cores = sorted_cores[0]
        ideal_speedup_values = [c / min_cores for c in sorted_cores]
        ax2.set_yticks(ideal_speedup_values)
        ax2.set_yticklabels([f'{val:.1f}' for val in ideal_speedup_values])
    
    # Set labels
    ax1.set_xlabel('Number of Cores', fontsize=12)
    ax1.set_ylabel('Computation Time (s)', fontsize=12, color='tab:blue')
    ax2.set_ylabel('Speedup', fontsize=12, color='tab:red')
    
    # Color the y-axis labels
    ax1.tick_params(axis='y', labelcolor='tab:blue')
    ax2.tick_params(axis='y', labelcolor='tab:red')
    
    # Add grid
    ax1.grid(True, alpha=0.3)
    
    # Combine legends from both axes
    lines1, labels1 = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(lines1 + lines2, labels1 + labels2, loc='upper center', fontsize=10)
    
    # Adjust layout
    plt.tight_layout()
    
    # Save figures
    pdf_path = pdf_dir / f'{test_case}.pdf'
    svg_path = svg_dir / f'{test_case}.svg'
    
    fig.savefig(pdf_path, format='pdf', bbox_inches='tight', dpi=300)
    fig.savefig(svg_path, format='svg', bbox_inches='tight')
    
    print(f"Saved {test_case}.pdf and {test_case}.svg")
    
    plt.close(fig)

print("All plots generated successfully!")

